# -*- coding: utf-8 -*-


from __future__ import print_function
import collections
import math
import numpy as np
import os
import random
import tensorflow as tf
import zipfile
from six.moves import range
from six.moves.urllib.request import urlretrieve
from sklearn.manifold import TSNE

# import matplotlib
# matplotlib.use('TkAgg')
from matplotlib import pylab

from common.datasets import Datasets
import os
import codecs
import logging


# dataset_dir = "~/Documents/ai_challenge/"
dataset_dir = "~/sentiment_analysis/AI-Challenge/"

dataset = Datasets(dataset_dir)

prj_path = os.path.dirname(__file__)

from collections import Counter


train_data = dataset.train_data()
dev_data = dataset.dev_data()
test_data_a = dataset.test_data_a()

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

logging.basicConfig(filename="log_att_raw.log", filemode="a", format="%(asctime)s-%(name)s-%(levelname)s-%(message)s", level=logging.INFO)

checkpointDir = "model/train_word2vector/"


def build_user_vocab_vector():
    # with open("./data/user_dict.txt", "r") as fp:
    #     all_lines = fp.readlines()
    #     all_lines = [_.strip() for _ in all_lines]
    #
    # all_lines_word_set = set(all_lines)

    import jieba

    jieba.load_userdict("data/user_dict.txt")
    jieba.enable_parallel()
    all_tencent_word_dict = dict()
    all_words = []

    for xxx in [train_data, dev_data, test_data_a]:
        for row in xxx.values:
            content = row[1]
            words_tmp = jieba.cut(content[1:-1], HMM=False)
            for w in words_tmp:
                # if not w.strip():
                #     continue
                # if w not in all_lines_word_set:
                # w_count = all_tencent_word_dict.get(w, 0)
                # w_count += 1
                # all_tencent_word_dict[w] = w_count

                all_words.append(w)

    # not_in_tencent_word_items = all_tencent_word_dict.items()
    # not_in_tencent_word_items_sort = sorted(not_in_tencent_word_items, key=lambda _: _[1], reverse=True)
    # return not_in_tencent_word_items_sort
    return all_words


def build_dataset(words):
    word_count = [['UNK', -1]]
    word_count.extend(collections.Counter(words).most_common(vocabulary_size - 1))
    dictionary = dict()
    for word, _ in word_count:
        dictionary[word] = len(dictionary)
    data = list()
    unk_count = 0
    for word in words:
        if word in dictionary:
            index = dictionary[word]
        else:
            index = 0  # dictionary['UNK']
            unk_count = unk_count + 1
        data.append(index)
    word_count[0][1] = unk_count
    reverse_dictionary = dict(zip(dictionary.values(), dictionary.keys()))
    return data, word_count, dictionary, reverse_dictionary


words = build_user_vocab_vector()

vocabulary_size = len(words)

data, count, dictionary, reverse_dictionary = build_dataset(words)
logging.info('Most common words (+UNK)', count[:5])
logging.info('Sample data', data[:10])
del words  # Hint to reduce memory.

data_index = 0


def generate_batch(batch_size, num_skips, skip_window):
    global data_index
    assert batch_size % num_skips == 0
    assert num_skips <= 2 * skip_window
    batch = np.ndarray(shape=(batch_size), dtype=np.int32)
    labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)
    span = 2 * skip_window + 1  # [ skip_window target skip_window ]
    buffer = collections.deque(maxlen=span)
    for _ in range(span):
        buffer.append(data[data_index])
        data_index = (data_index + 1) % len(data)
    for i in range(batch_size // num_skips):
        target = skip_window  # target label at the center of the buffer
        targets_to_avoid = [skip_window]
        for j in range(num_skips):
            while target in targets_to_avoid:
                target = random.randint(0, span - 1)
            targets_to_avoid.append(target)
            batch[i * num_skips + j] = buffer[skip_window]
            labels[i * num_skips + j, 0] = buffer[target]
        buffer.append(data[data_index])
        data_index = (data_index + 1) % len(data)
    return batch, labels


logging.info('data:', [reverse_dictionary[di] for di in data[:8]])

for num_skips, skip_window in [(2, 1), (4, 2)]:
    data_index = 0
    batch, labels = generate_batch(batch_size=8, num_skips=num_skips, skip_window=skip_window)
    logging.info('\nwith num_skips = %d and skip_window = %d:' % (num_skips, skip_window))
    logging.info('    batch:', [reverse_dictionary[bi] for bi in batch])
    logging.info('    labels:', [reverse_dictionary[li] for li in labels.reshape(8)])

batch_size = 32
embedding_size = 200  # Dimension of the embedding vector.
skip_window = 1  # How many words to consider left and right.
num_skips = 2  # How many times to reuse an input to generate a label.
# We pick a random validation set to sample nearest neighbors. here we limit the
# validation samples to the words that have a low numeric ID, which by
# construction are also the most frequent.
valid_size = 16  # Random set of words to evaluate similarity on.
valid_window = 100  # Only pick dev samples in the head of the distribution.
valid_examples = np.array(random.sample(range(valid_window), valid_size))
num_sampled = 64  # Number of negative examples to sample.

graph = tf.Graph()

with graph.as_default():
    # Input data.
    train_dataset = tf.placeholder(tf.int32, shape=[batch_size])
    train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])
    valid_dataset = tf.constant(valid_examples, dtype=tf.int32)

    # Variables.
    embeddings = tf.Variable(
        tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0))
    softmax_weights = tf.Variable(
        tf.truncated_normal([vocabulary_size, embedding_size],
                            stddev=1.0 / math.sqrt(embedding_size)))
    softmax_biases = tf.Variable(tf.zeros([vocabulary_size]))

    # Model.
    # Look up embeddings for inputs.
    embed = tf.nn.embedding_lookup(embeddings, train_dataset)
    # Compute the softmax loss, using a sample of the negative labels each time.
    loss = tf.reduce_mean(
        tf.nn.sampled_softmax_loss(weights=softmax_weights, biases=softmax_biases, inputs=embed,
                                   labels=train_labels, num_sampled=num_sampled, num_classes=vocabulary_size))

    # Optimizer.
    # Note: The optimizer will optimize the softmax_weights AND the embeddings.
    # This is because the embeddings are defined as a variable quantity and the
    # optimizer's `minimize` method will by default modify all variable quantities
    # that contribute to the tensor it is passed.
    # See docs on `tf.train.Optimizer.minimize()` for more details.
    optimizer = tf.train.AdagradOptimizer(1.0).minimize(loss)

    # Compute the similarity between minibatch examples and all embeddings.
    # We use the cosine distance:
    norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keepdims=True))
    normalized_embeddings = embeddings / norm
    valid_embeddings = tf.nn.embedding_lookup(
        normalized_embeddings, valid_dataset)
    similarity = tf.matmul(valid_embeddings, tf.transpose(normalized_embeddings))

num_steps = 100001

with tf.Session(graph=graph) as session:
    tf.global_variables_initializer().run()
    logging.info('Initialized')
    average_loss = 0
    for step in range(num_steps):
        batch_data, batch_labels = generate_batch(
            batch_size, num_skips, skip_window)
        feed_dict = {train_dataset: batch_data, train_labels: batch_labels}
        _, l = session.run([optimizer, loss], feed_dict=feed_dict)
        average_loss += l
        if step % 2000 == 0:
            if step > 0:
                average_loss = average_loss / 2000
            # The average loss is an estimate of the loss over the last 2000 batches.
            logging.info('Average loss at step %d: %f' % (step, average_loss))
            average_loss = 0
        # note that this is expensive (~20% slowdown if computed every 500 steps)
        if step % 10000 == 0:
            sim = similarity.eval()
            for i in range(valid_size):
                valid_word = reverse_dictionary[valid_examples[i]]
                top_k = 8  # number of nearest neighbors
                nearest = (-sim[i, :]).argsort()[1:top_k + 1]
                log = 'Nearest to %s:' % valid_word
                for k in range(top_k):
                    close_word = reverse_dictionary[nearest[k]]
                    log = '%s %s,' % (log, close_word)
                logging.info(log)
                if not tf.gfile.Exists(checkpointDir):
                    tf.gfile.MakeDirs(checkpointDir)
                saver = tf.train.Saver()
                saver.save(sess=session, save_path=checkpointDir + "model")

    final_embeddings = normalized_embeddings.eval()

try:
    with open("./data/final_embeddings.txt", "w+") as fppp:
        import json
        fppp.writelines(json.dumps(final_embeddings, ensure_ascii=False))
except Exception as e:
    logging.info(e)


# num_points = 400
#
# tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000, method='exact')
# two_d_embeddings = tsne.fit_transform(final_embeddings[1:num_points + 1, :])
#
#
# def plot(embeddings, labels):
#     assert embeddings.shape[0] >= len(labels), 'More labels than embeddings'
#
#     pylab.figure(figsize=(15, 15))  # in inches
#     for i, label in enumerate(labels):
#         x, y = embeddings[i, :]
#         pylab.scatter(x, y)
#         pylab.annotate(label, xy=(x, y), xytext=(5, 2), textcoords='offset points',
#                        ha='right', va='bottom')
#     pylab.show()
#
#
# words = [reverse_dictionary[i] for i in range(1, num_points + 1)]
# plot(two_d_embeddings, words)



